#include "base_model.h"


double mu_K_t = 6;
double cv_K_t = 0.3;
double sigma_K_t = mu_K_t * cv_K_t;
gsl_vector* K_t = gsl_vector_alloc(5);
gsl_vector* p_K_t = nullptr;

double mu_D_t = 5;
double sigma_D_t = 0.5;
gsl_vector* D_t = gsl_vector_alloc(10);
gsl_vector* p_D_t = nullptr;

double alpha = 0.95;
double p = 19;
double b = 19;
double c_e = 10;
double c_m = 2;
double h_e = 1;
double h_m = 0.2;
double theta = h_m / (1 - alpha);
int max_epoch = 2; // default value, consider T = 3

bool use_cache = false;
cache psi_1t_cache;
cache psi_2t_cache;
cache psi_3t_cache;

void init_variables(int T = 3, bool enable_cache = true) {
	max_epoch = T - 1;

	// K_t
	for (int i = 0; i < K_t->size; i++) {
		gsl_vector_set(K_t, i, 4 + i); // [4, 5, 6, 7, 8]
	}
	p_K_t = p_norm(K_t, mu_K_t, sigma_K_t);
	display(K_t);
	display(p_K_t);

	// D_t
	for (int i = 0; i < D_t->size; i++) {
		gsl_vector_set(D_t, i, i); // [0, 1, 2, ..., 9]
	}
	p_D_t = p_norm(D_t, mu_D_t, sigma_D_t);
	display(D_t);
	display(p_D_t);

	use_cache = enable_cache;
}

cache_data init_psi_1t_cache(double v_0, double v_1, double v_2, double v_3) {
	cache_data psi_1t_data;
	for (int k_i = 0; k_i < K_t->size; k_i++) {
		double K_t_i = gsl_vector_get(K_t, k_i);
		gsl_vector* e_range = range(0, v_1 - v_0 + 1, 1);
		gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);

		for (int d_i = 0; d_i < D_t->size; d_i++) {
			double D_t_i = gsl_vector_get(D_t, d_i);
			for (int e_i = 0; e_i < e_range->size; e_i++) {
				double e = gsl_vector_get(e_range, e_i);
				for (int z_i = 0; z_i < z_range->size; z_i++) {
					double z = gsl_vector_get(z_range, z_i);
					params p = make_tuple(v_1 - e - D_t_i, v_2 - e - D_t_i, v_3 - e - D_t_i, z - e - D_t_i);
					// psi_1t_data.insert(pair<params, double>(p, 0));
					psi_1t_data[p] = 0;
				}
			}
		}
		gsl_vector_free(e_range);
		gsl_vector_free(z_range);
	}
	return psi_1t_data;
}

cache_data init_psi_2t_cache(double v_0, double v_1, double v_2, double v_3)
{
	cache_data psi_2t_data;
	for (int k_i = 0; k_i < K_t->size; k_i++) {
		double K_t_i = gsl_vector_get(K_t, k_i);
		gsl_vector* y_range = range(v_1, v_2 + 1, 1);
		gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);

		for (int d_i = 0; d_i < D_t->size; d_i++) {
			double D_t_i = gsl_vector_get(D_t, d_i);
			for (int y_i = 0; y_i < y_range->size; y_i++) {
				double y = gsl_vector_get(y_range, y_i);
				for (int z_i = 0; z_i < z_range->size; z_i++) {
					double z = gsl_vector_get(z_range, z_i);
					params p = make_tuple(y - D_t_i, v_2 - D_t_i, v_3 - D_t_i, z - D_t_i);
					// psi_1t_data.insert(pair<params, double>(p, 0));
					psi_2t_data[p] = 0;
				}
			}
		}
		gsl_vector_free(y_range);
		gsl_vector_free(z_range);
	}
	return psi_2t_data;
}

cache_data init_psi_3t_cache(double v_0, double v_1, double v_2, double v_3)
{
	cache_data psi_3t_data;
	for (int k_i = 0; k_i < K_t->size; k_i++) {
		double K_t_i = gsl_vector_get(K_t, k_i);
		gsl_vector* y_range = range(v_2, v_3 + 1, 1);
		gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);

		for (int d_i = 0; d_i < D_t->size; d_i++) {
			double D_t_i = gsl_vector_get(D_t, d_i);
			for (int y_i = 0; y_i < y_range->size; y_i++) {
				double y = gsl_vector_get(y_range, y_i);
				for (int z_i = 0; z_i < z_range->size; z_i++) {
					double z = gsl_vector_get(z_range, z_i);
					params p = make_tuple(y - D_t_i, y - D_t_i, v_3 - D_t_i, z - D_t_i);
					// psi_1t_data.insert(pair<params, double>(p, 0));
					psi_3t_data[p] = 0;
				}
			}
		}
		gsl_vector_free(y_range);
		gsl_vector_free(z_range);
	}
	return psi_3t_data;
}

void copy_cache_data(cache_data* source, cache_data* target)
{
	for (cache_data::const_iterator it = source->begin(); it != source->end(); it++) {
		(*target)[it->first] = it->second;
	}
}

void init_model(double v_0, double v_1, double v_2, double v_3)
{
	int epoch = 1;
	psi_1t_cache[epoch] = init_psi_1t_cache(v_0, v_1, v_2, v_3);
	psi_2t_cache[epoch] = init_psi_2t_cache(v_0, v_1, v_2, v_3);
	psi_3t_cache[epoch] = init_psi_3t_cache(v_0, v_1, v_2, v_3);

	while (epoch != max_epoch) {
		cache_data c_1 = psi_1t_cache[epoch];
		cache_data c_2 = psi_2t_cache[epoch];
		cache_data c_3 = psi_3t_cache[epoch];
		// init next epoch cache data
		psi_1t_cache[epoch + 1] = cache_data();
		psi_2t_cache[epoch + 1] = cache_data();
		psi_3t_cache[epoch + 1] = cache_data();
		// begin iteration for psi_1t
		for (cache_data::const_iterator it = c_1.begin(); it != c_1.end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			cache_data c_psi_1t = init_psi_1t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_2t = init_psi_2t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_3t = init_psi_3t_cache(_v_0, _v_1, _v_2, _v_3);
			copy_cache_data(&c_psi_1t, &psi_1t_cache[epoch + 1]);
			copy_cache_data(&c_psi_2t, &psi_2t_cache[epoch + 1]);
			copy_cache_data(&c_psi_3t, &psi_3t_cache[epoch + 1]);
		}
		// begin iteration for psi_2t
		for (cache_data::const_iterator it = c_2.begin(); it != c_2.end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			cache_data c_psi_1t = init_psi_1t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_2t = init_psi_2t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_3t = init_psi_3t_cache(_v_0, _v_1, _v_2, _v_3);
			copy_cache_data(&c_psi_1t, &psi_1t_cache[epoch + 1]);
			copy_cache_data(&c_psi_2t, &psi_2t_cache[epoch + 1]);
			copy_cache_data(&c_psi_3t, &psi_3t_cache[epoch + 1]);
		}
		// begin iteration for psi_3t
		for (cache_data::const_iterator it = c_3.begin(); it != c_3.end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			cache_data c_psi_1t = init_psi_1t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_2t = init_psi_2t_cache(_v_0, _v_1, _v_2, _v_3);
			cache_data c_psi_3t = init_psi_3t_cache(_v_0, _v_1, _v_2, _v_3);
			copy_cache_data(&c_psi_1t, &psi_1t_cache[epoch + 1]);
			copy_cache_data(&c_psi_2t, &psi_2t_cache[epoch + 1]);
			copy_cache_data(&c_psi_3t, &psi_3t_cache[epoch + 1]);
		}

		epoch++;
	}

	cout << "model init success" << endl;
}

double init_calc(double v_0, double v_1, double v_2, double v_3) {
	int epoch = max_epoch;
	while (epoch != 1) {
		// calc cache of psi_1t
		for (cache_data::const_iterator it = psi_1t_cache[epoch - 1].begin();
			it != psi_1t_cache[epoch - 1].end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			psi_1t_cache[epoch - 1][it->first] = f_t(_v_0, _v_1, _v_2, _v_3, epoch);
		}
		// calc cache of psi_2t
		for (cache_data::const_iterator it = psi_2t_cache[epoch - 1].begin();
			it != psi_2t_cache[epoch - 1].end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			psi_2t_cache[epoch - 1][it->first] = f_t(_v_0, _v_1, _v_2, _v_3, epoch);
		}
		// calc cache of psi_3t
		for (cache_data::const_iterator it = psi_3t_cache[epoch - 1].begin();
			it != psi_3t_cache[epoch - 1].end(); it++) {
			double _v_0, _v_1, _v_2, _v_3;
			tie(_v_0, _v_1, _v_2, _v_3) = it->first;
			psi_3t_cache[epoch - 1][it->first] = f_t(_v_0, _v_1, _v_2, _v_3, epoch);
		}
		epoch--;
	}
	return f_t(v_0, v_1, v_2, v_3, 1);
}

void destroy_variables() {
	gsl_vector_free(K_t);
	gsl_vector_free(p_K_t);
	gsl_vector_free(D_t);
	gsl_vector_free(p_D_t);
}

double psi_1t(double v_1, double v_2, double v_3, int e, int z, int epoch)
{
	gsl_vector* sub1 = sub(v_1 - e, D_t);
	gsl_vector* add1 = add(D_t, -v_1 + e);
	gsl_vector* I1 = I(sub1);
	gsl_vector* I2 = I(add1);
	double val = h_e * E(I1, p_D_t)
		+ b * E(I2, p_D_t)
		+ c_e * (v_1 - e)
		+ c_m * (z - v_3)
		+ h_m * (v_3 - v_1)
		+ theta * e;
	if (epoch != max_epoch) {
		gsl_vector* _D_t = gsl_vector_alloc(D_t->size);
		for (int i = 0; i < D_t->size; i++) {
			double D_t_i = gsl_vector_get(D_t, i);
			double t = 0;
			if (use_cache) {
				params p = make_tuple(v_1 - e - D_t_i, v_2 - e - D_t_i, v_3 - e - D_t_i, z - e - D_t_i);
				t = psi_1t_cache[epoch][p];
			}
			else {
				t = f_t(v_1 - e - D_t_i, v_2 - e - D_t_i, v_3 - e - D_t_i, z - e - D_t_i, epoch + 1);
			}
			gsl_vector_set(_D_t, i, t);
		}
		val += alpha * E(_D_t, p_D_t);
		gsl_vector_free(_D_t);
	}
	gsl_vector_free(sub1);
	gsl_vector_free(add1);
	gsl_vector_free(I1);
	gsl_vector_free(I2);
	return val;
}

double psi_2t(double v_2, double v_3, int y, int z, int epoch)
{
	gsl_vector* sub1 = sub(y, D_t);
	gsl_vector* sub2 = add(D_t, y);
	gsl_vector* I1 = I(sub1);
	gsl_vector* I2 = I(sub2);
	double val = h_e * E(I1, p_D_t)
		+ b * E(I2, p_D_t)
		+ c_e * y
		+ c_m * (z - v_3)
		+ h_m * (v_3 - y);
	if (epoch != max_epoch) {
		gsl_vector* _D_t = gsl_vector_alloc(D_t->size);
		for (int i = 0; i < D_t->size; i++) {
			double D_t_i = gsl_vector_get(D_t, i);
			double t = 0;
			if (use_cache) {
				params p = make_tuple(y - D_t_i, v_2 - D_t_i, v_3 - D_t_i, z - D_t_i);
				t = psi_2t_cache[epoch][p];
			}
			else {
				t = f_t(y - D_t_i, v_2 - D_t_i, v_3 - D_t_i, z - D_t_i, epoch + 1);
			}
			gsl_vector_set(_D_t, i, t);
		}
		val += alpha * E(_D_t, p_D_t);
		gsl_vector_free(_D_t);
	}
	gsl_vector_free(sub1);
	gsl_vector_free(sub2);
	gsl_vector_free(I1);
	gsl_vector_free(I2);
	return val;
}

double psi_3t(double v_3, int y, int z, int epoch)
{
	gsl_vector* sub1 = sub(y, D_t);
	gsl_vector* sub2 = add(D_t, y);
	gsl_vector* I1 = I(sub1);
	gsl_vector* I2 = I(sub2);
	double val = h_e * E(I1, p_D_t)
		+ b * E(I2, p_D_t)
		+ c_e * y
		+ c_m * (z - v_3)
		+ h_m * (v_3 - y);
	if (epoch != max_epoch) {
		gsl_vector* _D_t = gsl_vector_alloc(D_t->size);
		for (int i = 0; i < D_t->size; i++) {
			double D_t_i = gsl_vector_get(D_t, i);
			double t = 0;
			if (use_cache) {
				params p = make_tuple(y - D_t_i, y - D_t_i, v_3 - D_t_i, z - D_t_i);
				t = psi_3t_cache[epoch][p];
			}
			else {
				t = f_t(y - D_t_i, y - D_t_i, v_3 - D_t_i, z - D_t_i, epoch + 1);
			}
			gsl_vector_set(_D_t, i, t);
		}
		val += alpha * E(_D_t, p_D_t);
		gsl_vector_free(_D_t);

	}
	gsl_vector_free(sub1);
	gsl_vector_free(sub2);
	gsl_vector_free(I1);
	gsl_vector_free(I2);
	return val;
}

double phi_1t(double v_0, double v_1, double v_2, double v_3, double K_t_i, int epoch)
{
	gsl_vector* e_range = range(0, v_1 - v_0 + 1, 1);
	gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);
	gsl_matrix* psi_1t_values = gsl_matrix_alloc(e_range->size, z_range->size);
	for (int i = 0; i < e_range->size; i++) {
		for (int j = 0; j < z_range->size; j++) {
			int e = gsl_vector_get(e_range, i);
			int z = gsl_vector_get(z_range, j);
			gsl_matrix_set(psi_1t_values, i, j, psi_1t(v_1, v_2, v_3, e, z, epoch));
		}
	}
	double m = gsl_matrix_min(psi_1t_values);
	gsl_vector_free(e_range);
	gsl_vector_free(z_range);
	gsl_matrix_free(psi_1t_values);
	return -c_e * v_0 + m;
}

double phi_2t(double v_0, double v_1, double v_2, double v_3, double K_t_i, int epoch)
{
	gsl_vector* y_range = range(v_1, v_2 + 1, 1);
	gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);
	gsl_matrix* psi_2t_values = gsl_matrix_alloc(y_range->size, z_range->size);
	for (int i = 0; i < y_range->size; i++) {
		for (int j = 0; j < z_range->size; j++) {
			int y = gsl_vector_get(y_range, i);
			int z = gsl_vector_get(z_range, j);
			gsl_matrix_set(psi_2t_values, i, j, psi_2t(v_2, v_3, y, z, epoch));
		}
	}
	double m = gsl_matrix_min(psi_2t_values);
	gsl_vector_free(y_range);
	gsl_vector_free(z_range);
	gsl_matrix_free(psi_2t_values);
	return -c_e * v_0 + m;
}

double phi_3t(double v_0, double v_1, double v_2, double v_3, double K_t_i, int epoch)
{
	gsl_vector* y_range = range(v_2, v_3 + 1, 1);
	gsl_vector* z_range = range(v_3, v_3 + K_t_i + 1, 1);
	gsl_matrix* psi_3t_values = gsl_matrix_alloc(y_range->size, z_range->size);
	for (int i = 0; i < y_range->size; i++) {
		for (int j = 0; j < z_range->size; j++) {
			int y = gsl_vector_get(y_range, i);
			int z = gsl_vector_get(z_range, j);
			gsl_matrix_set(psi_3t_values, i, j, psi_3t(v_3, y, z, epoch));
		}
	}
	double m = gsl_matrix_min(psi_3t_values);
	gsl_vector_free(y_range);
	gsl_vector_free(z_range);
	gsl_matrix_free(psi_3t_values);
	return -c_e * v_0 + m;
}

double f_t(double v_0, double v_1, double v_2, double v_3, int epoch)
{
	// cout << v_0 << v_1 << v_2 << v_3 << epoch << endl;
	gsl_vector* _K_t = gsl_vector_alloc(K_t->size);
	gsl_vector* m = gsl_vector_alloc(3);
	for (int i = 0; i < K_t->size; i++) {
		double K_t_i = gsl_vector_get(K_t, i);
		gsl_vector_set(m, 0, phi_1t(v_0, v_1, v_2, v_3, K_t_i, epoch));
		gsl_vector_set(m, 1, phi_2t(v_0, v_1, v_2, v_3, K_t_i, epoch));
		gsl_vector_set(m, 2, phi_3t(v_0, v_1, v_2, v_3, K_t_i, epoch));
		gsl_vector_set(_K_t, i, gsl_vector_min(m));
	}
	double val = E(_K_t, p_K_t);
	gsl_vector_free(_K_t);
	gsl_vector_free(m);
	return val;
}
